#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/X86Vector/X86VectorDialect.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Index/IR/IndexDialect.h"
#include "mlir/Dialect/Index/IR/IndexOps.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassRegistry.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/DialectConversion.h"
#include <iostream>

#include "crypto/Crypto.h"

namespace mlir {
#define GEN_PASS_DEF_CONVERTCRYPTOTOARM
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir


using namespace mlir;
using namespace vector;


//===----------------------------------------------------------------------===//
// Operation conversion
//===----------------------------------------------------------------------===//

namespace {

#define WORD 32
#define W_LEN 80
#define BLOCK_SIZE 64
#define STREAM_NUM 8
#define THRESHOLD 100000

class Sha1LoweringGpuPattern : public OpRewritePattern<crypto::Sha1multi_OP> {
public:
  using OpRewritePattern<crypto::Sha1multi_OP>::OpRewritePattern;

  LogicalResult matchAndRewrite(crypto::Sha1multi_OP op,
                                PatternRewriter &rewriter) const override {
    auto loc = op->getLoc();

    // get values needed from for-loop body for launching GPU kernel
    Value plaintext = op->getOperand(0);
    Value plaintext_length = op->getOperand(1);
    Value ciphertext = op->getOperand(2);
    Value num = op->getOperand(3);

    // prepare some vector types
    VectorType vectorI8x64 = VectorType::get({64}, rewriter.getI8Type());
    VectorType vectorI8x8 = VectorType::get({8}, rewriter.getI8Type());
    VectorType vectorI1x64 = VectorType::get({64}, rewriter.getI1Type());
    VectorType vectorI64x1 = VectorType::get({1}, rewriter.getI64Type());
    VectorType vectorI32x16 = VectorType::get({16}, rewriter.getI32Type());
    auto dynI8MemrefType = MemRefType::get({mlir::ShapedType::kDynamic}, rewriter.getI8Type());
    auto gpuAsyncTokenType = gpu::AsyncTokenType::get(rewriter.getContext());

    // index
    Value indexes[W_LEN];
    for (int i = 0; i < W_LEN; i++) {
        indexes[i] = rewriter.create<arith::ConstantIndexOp>(loc, i);
    }
    Value cn = rewriter.create<arith::IndexCastOp>(loc, /*out=*/rewriter.getIndexType(), /*in=*/plaintext_length);
    Value cN = rewriter.create<arith::IndexCastOp>(loc, /*out=*/rewriter.getIndexType(), /*in=*/num);

    // decide stream number
    Value streamNum = rewriter.create<arith::ConstantIndexOp>(loc, /*value=*/STREAM_NUM);
    Value threshold = rewriter.create<arith::ConstantIndexOp>(loc, /*value=*/THRESHOLD);
    streamNum = rewriter.create<arith::SelectOp>(loc, /*condition=*/rewriter.create<arith::CmpIOp>(loc, /*predicate=*/arith::CmpIPredicate::sle, /*lhs=*/cN, /*rhs=*/threshold), /*true_value=*/indexes[1], /*false_value=*/streamNum);
    Value streamNum_1 = rewriter.create<arith::ConstantIndexOp>(loc, /*value=*/STREAM_NUM - 1);
    // calculate the number of messages to be calculated by each stream
    Value segNum = rewriter.create<arith::DivUIOp>(loc, /*lhs=*/cN, /*rhs=*/streamNum);
    Value lastSegNum = rewriter.create<arith::SubIOp>(loc, /*lhs=*/cN, /*rhs=*/rewriter.create<arith::MulIOp>(loc, /*lhs=*/segNum, /*rhs=*/streamNum_1));
    // configure the kernel
    Value dynSharedMemSize = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(16384));
    Value blockSize = rewriter.create<arith::ConstantIndexOp>(loc, /*value=*/256);

    // multi stream
    rewriter.create<scf::ForOp>(
      loc, /*lowerBound=*/indexes[0], /*upperBound=*/streamNum, /*step=*/indexes[1], /*iterArgs=*/ValueRange{},
      /*odsArg4=*/[&](OpBuilder &builder, Location loc, Value ivs, ValueRange iargs){
        
        Value streamIdx = ivs;
        // shift
        Value msgNumShift = rewriter.create<arith::MulIOp>(loc, /*lhs=*/streamIdx, /*rhs=*/segNum);
        Value curMsgByteShift = rewriter.create<arith::MulIOp>(loc, /*lhs=*/msgNumShift, /*rhs=*/cn);
        Value curHashByteShift = rewriter.create<arith::MulIOp>(loc, /*lhs=*/msgNumShift, /*rhs=*/indexes[20]);
        // calculate the number of messages to be calculated by current stream
        Value curSegNum = rewriter.create<arith::SelectOp>(loc, /*condition=*/rewriter.create<arith::CmpIOp>(loc, /*predicate=*/arith::CmpIPredicate::eq, /*lhs=*/streamIdx, /*rhs=*/streamNum_1), /*true_value=*/lastSegNum, /*false_value=*/segNum);
        Value gridSize = rewriter.create<index::CeilDivUOp>(loc, /*lhs=*/curSegNum, /*rhs=*/blockSize);
        // calculate the size of messages to be calculated by current stream
        Value partMsgSize = rewriter.create<arith::MulIOp>(loc, /*lhs=*/curSegNum, /*rhs=*/cn);
        Value partHashSize = rewriter.create<arith::MulIOp>(loc, /*lhs=*/curSegNum, /*rhs=*/indexes[20]);
        Value partMsgMemHost = rewriter.create<memref::ViewOp>(loc, /*resultType0=*/dynI8MemrefType, /*source=*/plaintext, /*byte_shift=*/curMsgByteShift, /*sizes=*/ValueRange{partMsgSize});
        Value partHashMemHost = rewriter.create<memref::ViewOp>(loc, /*resultType0=*/dynI8MemrefType, /*source=*/ciphertext, /*byte_shift=*/curHashByteShift, /*sizes=*/ValueRange{partHashSize});

        // synchronize
        gpu::WaitOp waitOp = rewriter.create<gpu::WaitOp>(loc, /*asyncToken=*/gpuAsyncTokenType, /*asyncDependencies=*/ValueRange{});
        Value tokenWaitStart = waitOp.getResult(0);
        // allocate memory on device
        auto msgDevAllocOp = rewriter.create<gpu::AllocOp>(
          loc, 
          /*memref=*/dynI8MemrefType, 
          /*asyncToken=*/gpuAsyncTokenType, 
          /*asyncDependencies=*/ValueRange{tokenWaitStart}, 
          /*dynamicSizes=*/ValueRange{partMsgSize}, 
          /*symbolOperands=*/ValueRange{}
        );
        Value msgMemDev = msgDevAllocOp.getResult(0);
        Value tokenMsgAlloc = msgDevAllocOp.getResult(1);
        // asynchronously copy data from host to device
        auto memHost2DevOp = rewriter.create<gpu::MemcpyOp>(
          loc,
          /*asyncToken=*/gpuAsyncTokenType,
          /*asyncDependencies*/ValueRange{tokenMsgAlloc},
          /*dst=*/msgMemDev,
          /*src=*/partMsgMemHost
        );
        Value tokenMemHost2Dev = memHost2DevOp.getResult(0);
        auto hashDevAllocOp = rewriter.create<gpu::AllocOp>(
          loc,
          /*memref=*/dynI8MemrefType,
          /*asyncToken=*/gpuAsyncTokenType,
          /*asyncDependencies=*/ValueRange{tokenMemHost2Dev},
          /*dynamicSizes=*/ValueRange{partHashSize},
          /*symbolOperands=*/ValueRange{},
          /*hostShared=*/false
        );
        Value hashMemDev = hashDevAllocOp.getResult(0);
        Value tokenHashAlloc = hashDevAllocOp.getResult(1);

        // launch kernel
        gpu::LaunchOp launchOp = rewriter.create<gpu::LaunchOp>(
          loc, 
          /*gridSizeX=*/gridSize, /*gridSizeY=*/indexes[1], /*gridSizeZ=*/indexes[1],
          /*blockSizeX=*/blockSize, /*blockSizeY=*/indexes[1], /*blockSizeZ=*/indexes[1],
          /*dynamicSharedMemorySize*/dynSharedMemSize, /*asyncTokenType=*/gpuAsyncTokenType, /*asyncDependencies=*/ValueRange{tokenHashAlloc});
        Value tokenLaunch = launchOp.getResult(0);

        // Copy the hash results from device to host
        auto memDev2HostOp = rewriter.create<gpu::MemcpyOp>(loc, /*asyncToken=*/gpuAsyncTokenType, /*asyncDependencies*/ValueRange{tokenLaunch}, /*dst=*/partHashMemHost, /*src=*/hashMemDev);
        // deallcate the device memory
        auto msgDeallocOp = rewriter.create<gpu::DeallocOp>(loc, /*asyncToken=*/gpuAsyncTokenType, /*asyncDependencies=*/ValueRange{memDev2HostOp.getResult(0)}, /*memref=*/msgMemDev);
        auto hashDeallocOp = rewriter.create<gpu::DeallocOp>(loc, /*asyncToken=*/gpuAsyncTokenType, /*asyncDependencies=*/ValueRange{msgDeallocOp.getResult(0)}, /*memref=*/hashMemDev);
        // synchronize
        rewriter.create<gpu::WaitOp>(loc, /*asyncToken=*/std::nullopt, /*asyncDependencies=*/ValueRange{hashDeallocOp.getResult(0)});


        rewriter.create<scf::YieldOp>(loc);


        // Handwrite SHA-1 kenerl function body
        rewriter.setInsertionPointToStart(&launchOp.getBody().front());

        // Initialization
        Value H0 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(1732584193));  /*0x67452301*/
        Value H1 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(-271733879));  /*0xEFCDAB89*/
        Value H2 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(-1732584194)); /*0x98BADCFE*/
        Value H3 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(271733878));   /*0x10325476*/
        Value H4 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(-1009589776)); /*0xC3D2E1F0*/

        Value fullOneI32 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(-1));
        Value c0I8 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI8Type(), rewriter.getI8IntegerAttr(0));
        Value c_128I8 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI8Type(), rewriter.getI8IntegerAttr(-128));
        Value c1I32 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(1));
        Value c2I32 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(2));
        Value c5I32 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(5));
        Value c8I32 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(8));
        Value c8I64 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI64Type(), rewriter.getI64IntegerAttr(8));
        Value c27I32 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(27));
        Value c30I32 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(30));
        Value c31I32 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(31));
        Value c64I32 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(64));

        Value kernel_index[16];
        for (int i = 0; i < 16; i++) {
            kernel_index[i] = rewriter.create<arith::ConstantIndexOp>(loc, i);
        }
        Value c20Idx = rewriter.create<arith::ConstantIndexOp>(loc, 20);
        Value c56Idx = rewriter.create<arith::ConstantIndexOp>(loc, 56);
        Value c64Idx = rewriter.create<arith::ConstantIndexOp>(loc, 64);

        // calculate the unique id for each thread
        Value tIdX = rewriter.create<gpu::ThreadIdOp>(loc, /*dimension=*/::mlir::gpu::Dimension::x);
        Value blockIdxX = rewriter.create<gpu::BlockIdOp>(loc, /*dimension=*/::mlir::gpu::Dimension::x);
        Value tmpMul = rewriter.create<arith::MulIOp>(loc, /*lhs=*/blockIdxX, /*rhs=*/launchOp.getBlockSize().x);
        Value id = rewriter.create<arith::AddIOp>(loc, /*lhs=*/tmpMul, /*rhs=*/tIdX);
        // calculate stride
        Value totalStride = rewriter.create<arith::MulIOp>(loc, /*lhs=*/launchOp.getGridSize().x, /*rhs=*/launchOp.getBlockSize().x);

        // messages to be calculated by current thread
        rewriter.create<scf::ForOp>(
          loc, /*lowerBound=*/id, /*upperBound=*/curSegNum, /*step=*/totalStride, /*iterArgs=*/std::nullopt,
          /*odsArg4=*/[&](OpBuilder &builder, Location loc, Value ivs, ValueRange iargs){
            Value curMsgIdx = ivs;
            Value plaintext_length_idx = rewriter.create<arith::IndexCastOp>(loc, /*out=*/rewriter.getIndexType(), /*in=*/plaintext_length);
            // calculate the base address of current message
            Value curMsgBaseAddr = rewriter.create<arith::MulIOp>(loc, /*lhs=*/curMsgIdx, /*rhs=*/plaintext_length_idx);
            // calculate the base address of current to-be-stored ciphertext
            Value curHashBaseAddr = rewriter.create<arith::MulIOp>(loc, /*lhs=*/curMsgIdx, /*rhs=*/c20Idx);
            
            // calculate the padding size
            Value rem = rewriter.create<arith::RemSIOp>(loc, /*lhs=*/plaintext_length_idx, /*rhs=*/c64Idx);
            Value firstOneOffset = rem;
            Value diff = rewriter.create<arith::SubIOp>(loc, /*lhs=*/c64Idx, /*rhs=*/rem);
            Value isPadTwo = rewriter.create<arith::CmpIOp>(loc, /*predicate=*/arith::CmpIPredicate::sle, /*lhs=*/diff, /*rhs=*/kernel_index[8]);
            Value padSize1 = rewriter.create<arith::AddIOp>(loc, /*lhs=*/plaintext_length_idx, /*rhs=*/diff);
            Value padSize2 = rewriter.create<arith::AddIOp>(loc, /*lhs=*/padSize1, /*rhs=*/c64Idx);
            Value padSize = rewriter.create<arith::SelectOp>(loc, /*condition=*/isPadTwo, /*true_value=*/padSize2, /*false_value=*/padSize1);
            Value firstPaddingBlockIdx = rewriter.create<arith::DivUIOp>(loc, /*lhs=*/plaintext_length_idx, /*rhs=*/c64Idx);
            Value blockNum = rewriter.create<arith::DivUIOp>(loc, /*lhs=*/padSize, /*rhs=*/c64Idx);
            Value lastPaddingBlockIdx = rewriter.create<arith::SubIOp>(loc, /*lhs=*/blockNum, /*rhs=*/kernel_index[1]);

            scf::ForOp forOp = rewriter.create<scf::ForOp>(
              loc, /*lowerBound=*/kernel_index[0], /*upperBound=*/padSize, /*step=*/c64Idx, /*iterArgs=*/ValueRange{H0, H1, H2, H3, H4},
              /*odsArg4=*/[&](OpBuilder &builder, Location loc, Value ivs, ValueRange iargs){
                Value idx64 = ivs;
                Value a = iargs[0], b = iargs[1], c = iargs[2], d = iargs[3], e = iargs[4];
                Value H0_last = iargs[0], H1_last = iargs[1], H2_last = iargs[2], H3_last = iargs[3], H4_last = iargs[4];

                // decide whether or not to do the padding
                Value curBlockIdx = rewriter.create<arith::DivUIOp>(loc, /*lhs=*/idx64, /*rhs=*/c64Idx);
                Value isPaddingNow = rewriter.create<arith::CmpIOp>(loc, /*predicate=*/arith::CmpIPredicate::eq, /*lhs=*/curBlockIdx, /*rhs=*/firstPaddingBlockIdx);
                Value isLastBlock = rewriter.create<arith::CmpIOp>(loc, /*predicate=*/arith::CmpIPredicate::eq, /*lhs=*/curBlockIdx, /*rhs=*/lastPaddingBlockIdx);
                // calculate the address of read-in data
                Value curMsgAddr = rewriter.create<arith::AddIOp>(loc, /*lhs=*/curMsgBaseAddr, /*rhs=*/idx64);
                scf::IfOp ifPaddingOp = rewriter.create<scf::IfOp>(
                  loc, /*cond=*/isPaddingNow, /*thenBuilder=*/[&](OpBuilder &builder, Location loc) {
                    Value loadingMask = rewriter.create<vector::CreateMaskOp>(loc, /*resultType0=*/vectorI1x64, /*operand=*/ValueRange{firstOneOffset});
                    // load W[16] for each iteration
                    Value zeroBlock = rewriter.create<vector::SplatOp>(loc, /*aggregate=*/vectorI8x64, /*input=*/ValueRange{c0I8});
                    Value loadedData = rewriter.create<vector::MaskedLoadOp>(loc, /*result=*/vectorI8x64, /*base=*/msgMemDev, /*indices=*/curMsgAddr, /*mask=*/loadingMask, /*pass_thru=*/zeroBlock);
                    Value paddedBlock = rewriter.create<vector::InsertElementOp>(loc, /*source=*/c_128I8, /*dest=*/loadedData, /*position=*/firstOneOffset);
                    rewriter.create<scf::YieldOp>(loc, /*results=*/ValueRange{paddedBlock});
                  },
                  /*elseBuilder=*/[&](OpBuilder &builder, Location loc) {
                    Value zeroBlock = rewriter.create<vector::SplatOp>(loc, /*aggregate=*/vectorI8x64, /*input=*/ValueRange{c0I8});
                    Value loadedData = rewriter.create<vector::LoadOp>(loc, /*result=*/vectorI8x64, /*base=*/msgMemDev, /*indices=*/curMsgAddr);
                    Value result = rewriter.create<arith::SelectOp>(loc, /*condition=*/isLastBlock, /*true_value=*/zeroBlock, /*false_value=*/loadedData);
                    rewriter.create<scf::YieldOp>(loc, /*results=*/ValueRange{result});
                  });
                  scf::IfOp ifLastBlockOp = rewriter.create<scf::IfOp>(
                  loc, /*cond=*/isLastBlock, /*thenBuilder=*/[&](OpBuilder &builder, Location loc) {
                    Value extLen = rewriter.create<arith::ExtUIOp>(loc, /*out=*/rewriter.getI64Type(), /*in=*/plaintext_length);
                    Value msgBits = rewriter.create<arith::MulIOp>(loc, /*lhs=*/extLen, /*rhs=*/c8I64);
                    Value t0 = rewriter.create<vector::BroadcastOp>(loc, /*vector=*/vectorI64x1, /*source=*/msgBits);
                    Value t1 = rewriter.create<vector::BitCastOp>(loc, /*result=*/vectorI8x8, /*source=*/t0);
                    SmallVector<int64_t> mask = {7, 6, 5, 4, 3, 2, 1, 0};
                    Value t2 = rewriter.create<vector::ShuffleOp>(loc, /*v1=*/t1, /*v2=*/t1, /*odsArg2=*/mask);
                    scf::ForOp innerForOp = rewriter.create<scf::ForOp>(
                      loc, /*lowerBound=*/kernel_index[0], /*upperBound=*/kernel_index[8], /*step=*/kernel_index[1], /*iterArgs=*/ValueRange{ifPaddingOp.getResults()[0]},
                      /*odsArg4=*/[&](OpBuilder &builder, Location loc, Value ivs, ValueRange iargs) {
                        Value tailIdx = rewriter.create<arith::AddIOp>(loc, /*lhs=*/ivs, /*rhs=*/c56Idx);
                        Value t3 = rewriter.create<vector::ExtractElementOp>(loc, /*source=*/t2, /*position=*/ivs);
                        Value t4 = rewriter.create<vector::InsertElementOp>(loc, /*source=*/t3, /*dest=*/iargs[0], /*position=*/tailIdx);
                        rewriter.create<scf::YieldOp>(loc, /*results=*/ValueRange{t4});
                      });
                    rewriter.create<scf::YieldOp>(loc, /*results=*/ValueRange{innerForOp.getResults()[0]});
                  },
                  /*elseBuilder=*/[&](OpBuilder &builder, Location loc) {
                    rewriter.create<scf::YieldOp>(loc, /*results=*/ValueRange{ifPaddingOp.getResults()[0]});
                  });
                Value paddedMsg = ifLastBlockOp.getResults()[0];

                // calculating W[80]
                SmallVector<int64_t> inv64{3, 2, 1, 0, 7, 6, 5, 4, 11, 10, 9, 8, 15, 14, 13, 12, 19, 18, 17, 16, 23, 22, 21, 20, 27, 26, 25, 24, 31, 30, 29, 28, 35, 34, 33, 32, 39, 38, 37, 36, 43, 42, 41, 40, 47, 46, 45, 44, 51, 50, 49, 48, 55, 54, 53, 52, 59, 58, 57, 56, 63, 62, 61, 60};
                Value input0 = rewriter.create<vector::ShuffleOp>(loc, /*v1=*/paddedMsg, /*v2=*/paddedMsg, /*odsArg2=*/inv64);
                Value word = rewriter.create<vector::BitCastOp>(loc, vectorI32x16, input0);

                Value W[W_LEN];
                for (int i = 0; i < 16; i++) {
                    W[i] = rewriter.create<vector::ExtractElementOp>(loc, word, kernel_index[i]);
                }
                for (int i = 16; i < W_LEN; i++) {
                    Value xorTmp0 = rewriter.create<arith::XOrIOp>(loc, W[i - 3], W[i - 8]);
                    Value xorTmp1 = rewriter.create<arith::XOrIOp>(loc, xorTmp0, W[i - 14]);
                    Value xorRes = rewriter.create<arith::XOrIOp>(loc, xorTmp1, W[i - 16]);
                    Value leftShift = rewriter.create<arith::ShLIOp>(loc, xorRes, c1I32);
                    Value hbits = rewriter.create<arith::AndIOp>(loc, leftShift, fullOneI32);
                    Value lbits = rewriter.create<arith::ShRUIOp>(loc, xorRes, c31I32);
                    W[i] = rewriter.create<arith::OrIOp>(loc, hbits, lbits);
                }
                
                // sha1 round 1-20
                Value K0 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(1518500249));  /* 0x5a827999 */
                for (int i = 0; i < 20; i++) {
                    // ROTL(a, 5)
                    Value sli = rewriter.create<arith::ShLIOp>(loc, a, c5I32);
                    Value hbits = rewriter.create<arith::AndIOp>(loc, sli, fullOneI32);
                    Value lbits = rewriter.create<arith::ShRUIOp>(loc, a, c27I32);
                    Value tmp0 = rewriter.create<arith::OrIOp>(loc, hbits, lbits);
                    // f0(b, c, d)
                    Value andbc = rewriter.create<arith::AndIOp>(loc, b, c);
                    Value notb = rewriter.create<arith::XOrIOp>(loc, b, fullOneI32);
                    Value andbd = rewriter.create<arith::AndIOp>(loc, notb, d);
                    Value tmp1 = rewriter.create<arith::XOrIOp>(loc, andbc, andbd);
                    // T = ROTL(a, 5) + f0(b, c, d) + e + K0 + W[i]
                    Value t = rewriter.create<arith::AddIOp>(loc, tmp0, tmp1);
                    t = rewriter.create<arith::AddIOp>(loc, t, e);
                    t = rewriter.create<arith::AddIOp>(loc, t, K0);
                    t = rewriter.create<arith::AddIOp>(loc, t, W[i]);
                    // e = d
                    e = d;
                    // d = c
                    d = c;
                    // c = ROTL(b, 30)
                    sli = rewriter.create<arith::ShLIOp>(loc, b, c30I32);
                    hbits = rewriter.create<arith::AndIOp>(loc, sli, fullOneI32);
                    lbits = rewriter.create<arith::ShRUIOp>(loc, b, c2I32);
                    c = rewriter.create<arith::OrIOp>(loc, hbits, lbits);
                    // b = a
                    b = a;
                    a = t;
                }
                Value K1 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(1859775393));  /* 0x6ed9eba1 */
                // sha1 round 21-40
                for (int i = 20; i < 40; i++) {
                    // ROTL(a, 5)
                    Value sli = rewriter.create<arith::ShLIOp>(loc, a, c5I32);
                    Value hbits = rewriter.create<arith::AndIOp>(loc, sli, fullOneI32);
                    Value lbits = rewriter.create<arith::ShRUIOp>(loc, a, c27I32);
                    Value tmp0 = rewriter.create<arith::OrIOp>(loc, hbits, lbits);
                    // f1(b, c, d)
                    Value xorbc = rewriter.create<arith::XOrIOp>(loc, b, c);
                    Value tmp1 = rewriter.create<arith::XOrIOp>(loc, xorbc, d);
                    // T = ROTL(a, 5) + f1(b, c, d) + e + K1 + W[i]
                    Value t = rewriter.create<arith::AddIOp>(loc, tmp0, tmp1);
                    t = rewriter.create<arith::AddIOp>(loc, t, e);
                    t = rewriter.create<arith::AddIOp>(loc, t, K1);
                    t = rewriter.create<arith::AddIOp>(loc, t, W[i]);
                    // e = d
                    e = d;
                    // d = c
                    d = c;
                    // c = ROTL(b, 30)
                    sli = rewriter.create<arith::ShLIOp>(loc, b, c30I32);
                    hbits = rewriter.create<arith::AndIOp>(loc, sli, fullOneI32);
                    lbits = rewriter.create<arith::ShRUIOp>(loc, b, c2I32);
                    c = rewriter.create<arith::OrIOp>(loc, hbits, lbits);
                    // b = a
                    b = a;
                    a = t;
                }
                Value K2 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(-1894007588)); /* 0x8f1bbcdc */
                // sha1 round 41-60
                for (int i = 40; i < 60; i++) {
                    // ROTL(a, 5)
                    Value sli = rewriter.create<arith::ShLIOp>(loc, a, c5I32);
                    Value hbits = rewriter.create<arith::AndIOp>(loc, sli, fullOneI32);
                    Value lbits = rewriter.create<arith::ShRUIOp>(loc, a, c27I32);
                    Value tmp0 = rewriter.create<arith::OrIOp>(loc, hbits, lbits);
                    // f0(b, c, d)
                    Value andbc = rewriter.create<arith::AndIOp>(loc, b, c);
                    Value andbd = rewriter.create<arith::AndIOp>(loc, b, d);
                    Value andcd = rewriter.create<arith::AndIOp>(loc, c, d);
                    Value tmp1 = rewriter.create<arith::XOrIOp>(loc, andbc, andbd);
                    tmp1 = rewriter.create<arith::XOrIOp>(loc, tmp1, andcd);
                    // T = ROTL(a, 5) + f2(b, c, d) + e + K2 + W[i]
                    Value t = rewriter.create<arith::AddIOp>(loc, tmp0, tmp1);
                    t = rewriter.create<arith::AddIOp>(loc, t, e);
                    t = rewriter.create<arith::AddIOp>(loc, t, K2);
                    t = rewriter.create<arith::AddIOp>(loc, t, W[i]);
                    // e = d
                    e = d;
                    // d = c
                    d = c;
                    // c = ROTL(b, 30)
                    sli = rewriter.create<arith::ShLIOp>(loc, b, c30I32);
                    hbits = rewriter.create<arith::AndIOp>(loc, sli, fullOneI32);
                    lbits = rewriter.create<arith::ShRUIOp>(loc, b, c2I32);
                    c = rewriter.create<arith::OrIOp>(loc, hbits, lbits);
                    // b = a
                    b = a;
                    a = t;
                }
                Value K3 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(-899497514));  /* 0xca62c1d6 */
                // sha1 round 61-80
                for (int i = 60; i < 80; i++) {
                    // ROTL(a, 5)
                    Value sli = rewriter.create<arith::ShLIOp>(loc, a, c5I32);
                    Value hbits = rewriter.create<arith::AndIOp>(loc, sli, fullOneI32);
                    Value lbits = rewriter.create<arith::ShRUIOp>(loc, a, c27I32);
                    Value tmp0 = rewriter.create<arith::OrIOp>(loc, hbits, lbits);
                    // f1(b, c, d)
                    Value xorbc = rewriter.create<arith::XOrIOp>(loc, b, c);
                    Value tmp1 = rewriter.create<arith::XOrIOp>(loc, xorbc, d);
                    // T = ROTL(a, 5) + f3(b, c, d) + e + K3 + W[i]
                    Value t = rewriter.create<arith::AddIOp>(loc, tmp0, tmp1);
                    t = rewriter.create<arith::AddIOp>(loc, t, e);
                    t = rewriter.create<arith::AddIOp>(loc, t, K3);
                    t = rewriter.create<arith::AddIOp>(loc, t, W[i]);
                    // e = d
                    e = d;
                    // d = c
                    d = c;
                    // c = ROTL(b, 30)
                    sli = rewriter.create<arith::ShLIOp>(loc, b, c30I32);
                    hbits = rewriter.create<arith::AndIOp>(loc, sli, fullOneI32);
                    lbits = rewriter.create<arith::ShRUIOp>(loc, b, c2I32);
                    c = rewriter.create<arith::OrIOp>(loc, hbits, lbits);
                    // b = a
                    b = a;
                    a = t;
                }

                // H0[i] = H0[i-1] + a
                a = rewriter.create<arith::AddIOp>(loc, H0_last, a);
                // H1[i] = H1[i-1] + b
                b = rewriter.create<arith::AddIOp>(loc, H1_last, b);
                // H2[i] = H2[i-1] + c
                c = rewriter.create<arith::AddIOp>(loc, H2_last, c);
                // H3[i] = H3[i-1] + d
                d = rewriter.create<arith::AddIOp>(loc, H3_last, d);
                // H4[i] = H4[i-1] + e
                e = rewriter.create<arith::AddIOp>(loc, H4_last, e);

                builder.create<scf::YieldOp>(loc, ValueRange{a, b, c, d, e});
              }
            );

            // Save the final hash value
            H0 = forOp.getResults()[0];
            H1 = forOp.getResults()[1];
            H2 = forOp.getResults()[2];
            H3 = forOp.getResults()[3];
            H4 = forOp.getResults()[4];

            // TODO: single hash now
            VectorType vectorI8x20 = VectorType::get({20}, rewriter.getI8Type());
            VectorType vectorI32x5 = VectorType::get({5}, rewriter.getI32Type());
            // store the final hash value to vector
            Value abcde = rewriter.create<vector::SplatOp>(loc, H4, vectorI32x5);
            abcde = rewriter.create<vector::InsertOp>(loc, /*source=*/H0, /*dest=*/abcde, /*position=*/indexes[0]);
            abcde = rewriter.create<vector::InsertOp>(loc, /*source=*/H1, /*dest=*/abcde, /*position=*/indexes[1]);
            abcde = rewriter.create<vector::InsertOp>(loc, /*source=*/H2, /*dest=*/abcde, /*position=*/indexes[2]);
            abcde = rewriter.create<vector::InsertOp>(loc, /*source=*/H3, /*dest=*/abcde, /*position=*/indexes[3]);
            // store the final hash value from vector to memref
            SmallVector<int64_t> inv20{3, 2, 1, 0, 7, 6, 5, 4, 11, 10, 9, 8, 15, 14, 13, 12, 19, 18, 17, 16};
            abcde = rewriter.create<vector::BitCastOp>(loc, vectorI8x20, abcde);
            abcde = rewriter.create<vector::ShuffleOp>(loc, abcde, abcde, inv20);
            rewriter.create<vector::StoreOp>(loc, abcde, hashMemDev, curHashBaseAddr);

            rewriter.create<scf::YieldOp>(loc);
          });

        // Insert the `gpu.terminate` op.
        rewriter.setInsertionPointToEnd(&launchOp.getBody().back());
        rewriter.create<gpu::TerminatorOp>(loc);
          
      });
    
    // rewriter.setInsertionPointAfter(streamForOp);
    // rewriter.create<gpu::WaitOp>(loc, /*asyncToken=*/gpuAsyncTokenType, /*asyncDependencies=*/ValueRange{});

    rewriter.eraseOp(op);
    return success();
  }
};

} // namespace


//===----------------------------------------------------------------------===//
// MD5
//===----------------------------------------------------------------------===//

namespace {

#define X_LEN 16
#define BLOCK_SIZE 64
#define STREAM_NUM 8
#define THRESHOLD 100000

class MD5LoweringGpuPattern : public OpRewritePattern<crypto::MD5multiOP> {
public:
  using OpRewritePattern<crypto::MD5multiOP>::OpRewritePattern;

  LogicalResult matchAndRewrite(crypto::MD5multiOP op,
                                PatternRewriter &rewriter) const override {
    auto loc = op->getLoc();

    // get values needed from for-loop body for launching GPU kernel
    Value plaintext = op->getOperand(0);
    Value plainLen = op->getOperand(1);
    Value ciphertext = op->getOperand(2);
    Value num = op->getOperand(3);

    // prepare some necessary types
    auto dynI8MemrefType = MemRefType::get({mlir::ShapedType::kDynamic}, rewriter.getI8Type());
    auto gpuAsyncTokenType = gpu::AsyncTokenType::get(rewriter.getContext());
    // prepare some necessary constant values
    Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
    Value c1 = rewriter.create<arith::ConstantIndexOp>(loc, 1);
    Value c16 = rewriter.create<arith::ConstantIndexOp>(loc, 16);
    Value cn = rewriter.create<arith::IndexCastOp>(loc, /*out=*/rewriter.getIndexType(), /*in=*/plainLen);
    Value cN = rewriter.create<arith::IndexCastOp>(loc, /*out=*/rewriter.getIndexType(), /*in=*/num);

    // decide stream number
    Value streamNum = rewriter.create<arith::ConstantIndexOp>(loc, /*value=*/STREAM_NUM);
    Value threshold = rewriter.create<arith::ConstantIndexOp>(loc, /*value=*/THRESHOLD);
    streamNum = rewriter.create<arith::SelectOp>(loc, /*condition=*/rewriter.create<arith::CmpIOp>(loc, /*predicate=*/arith::CmpIPredicate::sle, /*lhs=*/cN, /*rhs=*/threshold), /*true_value=*/c1, /*false_value=*/streamNum);
    Value streamNum_1 = rewriter.create<arith::ConstantIndexOp>(loc, /*value=*/STREAM_NUM - 1);
    // calculate the number of messages to be calculated by each stream
    Value segNum = rewriter.create<arith::DivUIOp>(loc, /*lhs=*/cN, /*rhs=*/streamNum);
    Value lastSegNum = rewriter.create<arith::SubIOp>(loc, /*lhs=*/cN, /*rhs=*/rewriter.create<arith::MulIOp>(loc, /*lhs=*/segNum, /*rhs=*/streamNum_1));
    // configure the kernel
    Value dynSharedMemSize = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(16384));
    Value blockSize = rewriter.create<arith::ConstantIndexOp>(loc, /*value=*/512);

    // multi stream
    rewriter.create<scf::ForOp>(
      loc, /*lowerBound=*/c0, /*upperBound=*/streamNum, /*step=*/c1, /*iterArgs=*/ValueRange{},
      /*odsArg4=*/[&](OpBuilder &builder, Location loc, Value ivs, ValueRange iargs){
        Value streamIdx = ivs;

        // shift
        Value msgNumShift = rewriter.create<arith::MulIOp>(loc, /*lhs=*/streamIdx, /*rhs=*/segNum);
        Value curMsgByteShift = rewriter.create<arith::MulIOp>(loc, /*lhs=*/msgNumShift, /*rhs=*/cn);
        Value curHashByteShift = rewriter.create<arith::MulIOp>(loc, /*lhs=*/msgNumShift, /*rhs=*/c16);
        // calculate the number of messages to be calculated by current stream
        Value curSegNum = rewriter.create<arith::SelectOp>(loc, /*condition=*/rewriter.create<arith::CmpIOp>(loc, /*predicate=*/arith::CmpIPredicate::eq, /*lhs=*/streamIdx, /*rhs=*/streamNum_1), /*true_value=*/lastSegNum, /*false_value=*/segNum);
        Value gridSize = rewriter.create<arith::CeilDivUIOp>(loc, /*lhs=*/curSegNum, /*rhs=*/blockSize);
        // calculate the size of messages to be calculated by current stream
        Value partMsgSize = rewriter.create<arith::MulIOp>(loc, /*lhs=*/curSegNum, /*rhs=*/cn);
        Value partHashSize = rewriter.create<arith::MulIOp>(loc, /*lhs=*/curSegNum, /*rhs=*/c16);
        Value partMsgMemHost = rewriter.create<memref::ViewOp>(loc, /*resultType0=*/dynI8MemrefType, /*source=*/plaintext, /*byte_shift=*/curMsgByteShift, /*sizes=*/ValueRange{partMsgSize});
        Value partHashMemHost = rewriter.create<memref::ViewOp>(loc, /*resultType0=*/dynI8MemrefType, /*source=*/ciphertext, /*byte_shift=*/curHashByteShift, /*sizes=*/ValueRange{partHashSize});

        // synchronize
        gpu::WaitOp waitOp = rewriter.create<gpu::WaitOp>(loc, /*asyncToken=*/gpuAsyncTokenType, /*asyncDependencies=*/ValueRange{});
        Value tokenWaitStart = waitOp.getResult(0);
        // allocate memory on device
        auto msgDevAllocOp = rewriter.create<gpu::AllocOp>(
          loc, 
          /*memref=*/dynI8MemrefType, 
          /*asyncToken=*/gpuAsyncTokenType, 
          /*asyncDependencies=*/ValueRange{tokenWaitStart}, 
          /*dynamicSizes=*/ValueRange{partMsgSize}, 
          /*symbolOperands=*/ValueRange{}
        );
        Value msgMemDev = msgDevAllocOp.getResult(0);
        Value tokenMsgAlloc = msgDevAllocOp.getResult(1);
        // asynchronously copy data from host to device
        auto memHost2DevOp = rewriter.create<gpu::MemcpyOp>(
          loc,
          /*asyncToken=*/gpuAsyncTokenType,
          /*asyncDependencies*/ValueRange{tokenMsgAlloc},
          /*dst=*/msgMemDev,
          /*src=*/partMsgMemHost
        );
        Value tokenMemHost2Dev = memHost2DevOp.getResult(0);
        auto hashDevAllocOp = rewriter.create<gpu::AllocOp>(
          loc,
          /*memref=*/dynI8MemrefType,
          /*asyncToken=*/gpuAsyncTokenType,
          /*asyncDependencies=*/ValueRange{tokenMemHost2Dev},
          /*dynamicSizes=*/ValueRange{partHashSize},
          /*symbolOperands=*/ValueRange{},
          /*hostShared=*/false
        );
        Value hashMemDev = hashDevAllocOp.getResult(0);
        Value tokenHashAlloc = hashDevAllocOp.getResult(1);


        // launch kernel
        gpu::LaunchOp launchOp = rewriter.create<gpu::LaunchOp>(
          loc, 
          /*gridSizeX=*/gridSize, /*gridSizeY=*/c1, /*gridSizeZ=*/c1,
          /*blockSizeX=*/blockSize, /*blockSizeY=*/c1, /*blockSizeZ=*/c1,
          /*dynamicSharedMemorySize*/dynSharedMemSize, /*asyncTokenType=*/gpuAsyncTokenType, /*asyncDependencies=*/ValueRange{tokenHashAlloc});
        Value tokenLaunch = launchOp.getResult(0);

        // Copy the hash results from device to host
        auto memDev2HostOp = rewriter.create<gpu::MemcpyOp>(loc, /*asyncToken=*/gpuAsyncTokenType, /*asyncDependencies*/ValueRange{tokenLaunch}, /*dst=*/partHashMemHost, /*src=*/hashMemDev);
        // deallocate the device memory
        auto msgDeallocOp = rewriter.create<gpu::DeallocOp>(loc, /*asyncToken=*/gpuAsyncTokenType, /*asyncDependencies=*/ValueRange{memDev2HostOp.getResult(0)}, /*memref=*/msgMemDev);
        auto hashDeallocOp = rewriter.create<gpu::DeallocOp>(loc, /*asyncToken=*/gpuAsyncTokenType, /*asyncDependencies=*/ValueRange{msgDeallocOp.getResult(0)}, /*memref=*/hashMemDev);
        // synchronize
        rewriter.create<gpu::WaitOp>(loc, /*asyncToken=*/std::nullopt, /*asyncDependencies=*/ValueRange{hashDeallocOp.getResult(0)});

        // end of stream
        rewriter.create<scf::YieldOp>(loc);


        // Handwrite MD5 kenerl function body
        rewriter.setInsertionPointToStart(&launchOp.getBody().front());


        // calculate the unique id for each thread
        Value tIdX = rewriter.create<gpu::ThreadIdOp>(loc, /*dimension=*/::mlir::gpu::Dimension::x);
        Value blockIdxX = rewriter.create<gpu::BlockIdOp>(loc, /*dimension=*/::mlir::gpu::Dimension::x);
        Value tmpMul = rewriter.create<arith::MulIOp>(loc, /*lhs=*/blockIdxX, /*rhs=*/launchOp.getBlockSize().x);
        Value id = rewriter.create<arith::AddIOp>(loc, /*lhs=*/tmpMul, /*rhs=*/tIdX);
        // calculate stride
        Value totalStride = rewriter.create<arith::MulIOp>(loc, /*lhs=*/launchOp.getGridSize().x, /*rhs=*/launchOp.getBlockSize().x);

        // messages to be calculated by current thread
        rewriter.create<scf::ForOp>(
          loc, /*lowerBound=*/id, /*upperBound=*/segNum, /*step=*/totalStride, /*iterArgs=*/std::nullopt,
          /*odsArg4=*/[&](OpBuilder &builder, Location loc, Value ivs, ValueRange iargs){
            Value curMsgIdx = ivs;
            // calculate the base address of current message
            Value curMsgBaseAddr = rewriter.create<arith::MulIOp>(loc, /*lhs=*/curMsgIdx, /*rhs=*/cn);
            // calculate the base address of current to-be-stored ciphertext
            Value curHashBaseAddr = rewriter.create<arith::MulIOp>(loc, /*lhs=*/curMsgIdx, /*rhs=*/c16);

            Value input = rewriter.create<memref::ViewOp>(loc, /*resultType0=*/dynI8MemrefType, /*source=*/msgMemDev, /*byte_shift=*/curMsgBaseAddr, /*sizes=*/cn);
            Value output = rewriter.create<memref::ViewOp>(loc, /*resultType0=*/dynI8MemrefType, /*source=*/hashMemDev, /*byte_shift=*/curHashBaseAddr, /*sizes=*/c16);

            rewriter.create<crypto::MD5OP>(loc, /*input=*/input, /*msgLen=*/plainLen, /*output=*/output);

            rewriter.create<scf::YieldOp>(loc);
          });

        // Insert the `gpu.terminate` op.
        rewriter.setInsertionPointToEnd(&launchOp.getBody().back());
        rewriter.create<gpu::TerminatorOp>(loc);

      });

    rewriter.eraseOp(op);
    return success();
  }

};

} // namespace


void populateConvertCryptoToGpuPatterns(RewritePatternSet &patterns) {
  patterns.add<Sha1LoweringGpuPattern>(patterns.getContext());
  patterns.add<MD5LoweringGpuPattern>(patterns.getContext());
}

namespace {
class LowerCryptoToGpuPass : public PassWrapper<LowerCryptoToGpuPass, OperationPass<ModuleOp>> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerCryptoToGpuPass)
  LowerCryptoToGpuPass() = default;
  LowerCryptoToGpuPass(const LowerCryptoToGpuPass &) {}

  StringRef getArgument() const final { return "lower-crypto-to-gpu"; }
  StringRef getDescription() const final { return "Lower Crypto Dialect to GPU Dialect."; }

  void runOnOperation() override;

  void getDependentDialects(DialectRegistry &registry) const override {
    registry
        .insert<func::FuncDialect,
                mlir::LLVM::LLVMDialect,
                memref::MemRefDialect,
                vector::VectorDialect,
                arith::ArithDialect,
                affine::AffineDialect,
                scf::SCFDialect,
                gpu::GPUDialect,
                index::IndexDialect
                >();
  }
};
} // end anonymous namespace.

void LowerCryptoToGpuPass::runOnOperation() {
  MLIRContext *context = &getContext();
  ModuleOp module = getOperation();

  ConversionTarget target(*context);
  target.addLegalDialect<
  func::FuncDialect,
  vector::VectorDialect,
  mlir::LLVM::LLVMDialect,
  memref::MemRefDialect,
  arith::ArithDialect,
  affine::AffineDialect,
  scf::SCFDialect,
  gpu::GPUDialect,
  index::IndexDialect
  >();
  target.addLegalOp<
  ModuleOp,
  func::FuncOp,
  func::ReturnOp,
  gpu::GPUFuncOp
  >();

  RewritePatternSet patterns(context);
  populateConvertCryptoToGpuPatterns(patterns);

  if (failed(applyPartialConversion(module, target, std::move(patterns))))
    signalPassFailure();
}

void registerLowerCryptoToGpuPass() { PassRegistration<LowerCryptoToGpuPass>(); }

std::unique_ptr<Pass> createLowerCryptoToGpuPass(){ return std::make_unique<LowerCryptoToGpuPass>(); }
